[ROCm][Perf] Split MiniMax-M3 prefill index-score K loop - #52059

Note: Validation required cherry-picking two fixes on top of the nightly image:

• PR #51585 — Fix CUDA graph capture failure (zeroed CPU query offsets broke mixed-attention metadata builders)
• PR #51632 — Fix Triton fused shared expert alignment (incorrect expert count caused silent accuracy loss)

Base image: vllm/vllm-openai-rocm:nightly-3ee2df30337a301164c46ae444b76ee67e71c106

Purpose

MiniMax-M3 prefill computes index scores by scanning every visible 128-token index-K block for each query block.

The current gfx942 kernel uses 64 query tokens per program and processes the full K range in one loop. At long context lengths, this repeatedly loads index-K data and leaves insufficient parallel work for the GPU.

Suggested Fix

For gfx942 prefill with at least 128 query tokens and an 8K sequence length:

• Increase the query block size from 64 to 128.
• Split the state-free K loop across multiple programs.
• Select a power-of-two split count using the query workload, sequence length, and available compute units.
• Cap the split count at 16.

Each split writes a separate range of score columns, so no atomics or reduction pass are required. Other architectures and smaller shapes retain the existing 64-token, single-split configuration.

Test Plan

pytest -q tests/kernels/attention/test_minimax_m3.py \
  -k prefill_index_topk_correctness

The test covers the existing short-context path and the gfx942 long-context split path. It compares every score written by the kernel, including causal masking, and verifies the final top-k block sets.

export HIP_VISIBLE_DEVICES=0,1,2,3
export VLLM_ROCM_USE_AITER=1
export VLLM_USE_BREAKABLE_CUDAGRAPH=0
export VLLM_ROCM_USE_AITER_MOE=0
export VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1

vllm serve EmbeddedLLM/MiniMax-M3-FP8-dynamic \
  --served-model-name minimax-m3 \
  --tensor-parallel-size 4 \
  --block-size 128 \
  --max-model-len 262144 \
  --gpu-memory-utilization 0.92 \
  --enable-chunked-prefill \
  --max-num-batched-tokens 32768 \
  --max-num-seqs 32 \
  --no-enable-prefix-caching \
  --async-scheduling \
  --tool-call-parser minimax_m3 \
  --enable-auto-tool-choice \
  --reasoning-parser minimax_m3 \
  --kv-cache-dtype fp8 \
  --language-model-only \
  --port 8000 \
  --attention-backend ROCM_AITER_UNIFIED_ATTN

Test Result

Kernel accuracy: 2/2 passed (short-context fallback + long-context split-K).

End-to-end serving on 4x MI325X (gfx942), TP4, EmbeddedLLM/MiniMax-M3-FP8-dynamic, FP8 KV cache:

Long-context prefill improvement scales with sequence length.

Short-context workloads are unaffected (selector falls back to unsplit path). Long-context prefill shows a consistent TTFT improvement.